from new_editor.search_agent import model
from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.get_json import get_json
import time
from random import randint


prompt_from_outline = ChatPromptTemplate.from_template("""
你负责根据要求撰写对应的段落内容(要求至少30字)
 - 综合标题{title}
 - 对应段落标题以及内容{sections}
 - 你当前需要完成的第一个段落
""")

prompt_from_check = ChatPromptTemplate.from_template("""
你负责根据校对员指出的不真实的信息意见修改编撰的段落内容(要求30字左右)
 - 你之前撰写的内容{content}
 - 校对员给出的指导意见{ideas}
 - 注意,此时你只负责修正不真实的部分,不必大动篇幅修改
""")

prompt_from_comment = ChatPromptTemplate.from_template("""
你负责根据评论员指出的意见修改编撰的段落内容(要求30字左右)
 - 你之前撰写的内容{content}
 - 评论员给出的指导意见{ideas}
 - 注意,此时你只负责按照评论员给出意见进行修改
""")


def generate_content(state: EditorState):
    who = state['who']
    title = state['title']
    outlines = state['outlines']
    search = state['search']
    paragraphs = state['paragraphs']
    sections = outlines['sections']
    check_idea = state['check_idea']
    comment_idea = state['comment_idea']
    write_submit = state['write_submit']

    print("============ Write Agent ============")

    time.sleep(randint(7,9))


    response = ""
    if who == "outline":
        chain_outline = prompt_from_outline | model
        response = chain_outline.invoke({"title": title, "sections": sections})
        print(response.content)
    elif who == "check":
        chain_outline = prompt_from_check | model
        response = chain_outline.invoke({"content": write_submit, "ideas": check_idea['idea']})
        print(response.content)
    elif who == "comment":
        chain_outline = prompt_from_comment | model
        response = chain_outline.invoke({"content": write_submit, "ideas": comment_idea['idea']})
        print(response.content)

    return {
        "who": "write",
        "title": title,  # 制定标题
        "search": search,  # 搜集信息
        "outlines": outlines,  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 段落信息
        "write_submit": response.content,
        "check_idea": {},  # 校对意见
        "comment_idea": {}  # 评论意见
    }


# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
